from Agent import *
from prompts import *
from split_function import *
from extract_component import *

# from add_stest_functions import * 
# from remove_stest_functions import *
from remove_alloc_test_blocks import *
from process_test_functions import remove_stest_functions, add_stest_functions, add_test_annotation, remove_out_of_memory_test
from process_test_functions import comment_code, remove_main_and_tests
from first_round_optimize import first_round_optimize
from process_crate import add_crate_and_lib, add_use_crate, add_alloc_set_limit, add_print_num_assert_fun
from process_crate import remove_extern_declaration_for_srcfile, remove_extern_declaration_for_testfile
from process_crate import remove_struct_and_type_declaration, remove_c_label_blocks, remove_atomic_type, remove_no_mangle_annotation
from global_config import *
from split_function import get_dependency_ordered_functions
import subprocess
import re
import logging
import shutil
import time
import datetime
import select


# agent = Agent(
#     role="代码优化",
#     prompt="You are a proficient C and Rust advanced developer.",
#     temperature=MODEL_TEMP,
#     top_p=TOP_P
# )

deepseek_agent = DeepSeekAgent(
    role="DeepSeek",
    prompt="You are a proficient C and Rust advanced developer.",
    temperature=MODEL_TEMP,
    top_p=TOP_P
)
doubao_agent = DoubaoAgent(
    role="Doubao",
    prompt="You are a proficient C and Rust advanced developer.",
    temperature=MODEL_TEMP,
    top_p=TOP_P
)


token_usage = (0, 0, 0)

def setup_logging():
    LOG_PATH = 'Output/logs'
    os.makedirs(LOG_PATH, exist_ok=True)
    
    # 配置日志格式
    log_format = '%(asctime)s - %(levelname)s - %(message)s'
    date_format = '%Y-%m-%d %H:%M:%S'
    
    # 创建日志文件名（使用时间戳）
    log_file = f'{LOG_PATH}/conversion_{time.strftime("%Y%m%d_%H%M%S")}.log'
    
    logging.basicConfig(
        level=logging.INFO,
        format=log_format,
        datefmt=date_format,
        handlers=[
            logging.FileHandler(log_file, encoding='utf-8'),  # 文件处理器
            logging.StreamHandler()  # 控制台处理器
        ]
    )


def check_cur_path():
    current_directory = os.getcwd()
    if os.path.basename(current_directory)=="tool":
        return True
    return False

def cargo_test(timeout=TEST_TIMEOUT, test_target=None, multi_thread=False):
    current_directory = os.getcwd()
    os.chdir(f"{PROJ_PATH}")
    # 运行 cargo test 命令
    env = os.environ.copy()
    env["RUSTFLAGS"] = "-Awarnings"
    command = ["cargo", "test"]
    if test_target:
        command.append("--test")
        command.append(test_target)
    if not multi_thread:
        command.append("--")
        command.append("--test-threads=1")
    logging.info(f"command: {command}")

    # command = ["cargo", "test", "--", "--test-threads=1"]
    # 第一次构建，可能有crates更新   
    try:
        result = subprocess.run(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            env=env,
            timeout=timeout
        )
        os.chdir(current_directory)
        if result.returncode == 0:
            logging.info("所有测试通过！")
            return (True, "")
        else:
            logging.info("测试未通过！")
            error_msg = result.stderr
            logging.info(error_msg)
            return (False, error_msg)
    except subprocess.TimeoutExpired:
        logging.info("超时")
        if timeout == FIRST_TEST_TIMEOUT:
            logging.info("可能是cargo crates未能更新，请在项目路径下手动运行cargo update")
        os.chdir(current_directory)
        return (False, "Timeout error, Please ensure that your program does not cause errors or infinite loops.")

# 从```中提取rust代码
# 返回代码/""
def extract_rust_code(review_text: str) -> str:
    if review_text is None or review_text == "":
        return ""
    code_blocks = re.findall(r'```rust\n(.*?)```', review_text, re.DOTALL)
    if code_blocks:
        review_text = code_blocks[0].strip()
    review_text = re.sub(r'//.*$', '', review_text, flags=re.MULTILINE)
    review_text = re.sub(r'[\u4e00-\u9fff]+', '', review_text)
    return review_text.strip()

# 将文件中的旧代码替换为新代码
# 如果找不到目标段落，返回F
def function_replace(file_path, origin_code, new_code):
    with open(file_path,'r') as file:
        content = file.read()
    if origin_code not in content:
        logging.info(f"下列原函数未能在文件{file_path}中查找到：\n{origin_code}")
        return False
    with open(file_path,'w') as file:
        file.write(content.replace(origin_code, new_code))
    return True

# 将文件中的新代码还原为旧代码
# 如果找不到目标段落，返回F
def function_restore(file_path, origin_code, new_code):
    with open(file_path,'r') as file:
        content = file.read()
    if new_code not in content:
        logging.info(f"下列新函数未能在文件{file_path}中查找到：\n{new_code}")
        return False
    with open(file_path,'w') as file:
        file.write(content.replace(new_code, origin_code))
    return True

def append_content(file_path, content):
    with open(file_path, 'a') as file:
        file.write(content)

def remove_content(file_path, content):
    with open(file_path, 'r') as file:
        origin_content = file.read()
    origin_content = origin_content.replace(content, "")
    with open(file_path, 'w') as file:
        file.write(origin_content)

# 对指定的代码（origin_code）迭代优化
# 需传入prompt（包括输入输出示例和user_ask）
# 返回T/F
def optimize(file_path, origin_code, user_ask, agent, example_input=[], example_output=[], opt_round=-1, multi_thread=False, test_target=None, assert_count=None):
    cnt = 0
    flag = False
    cur_user_ask = user_ask
    while cnt < MAX_ATTEMPS:
        logging.info("prompt输入：")
        logging.info(cur_user_ask)
        output, usage = agent.generate_response(user_input=cur_user_ask, example_input=example_input, example_output=example_output)
        output = extract_rust_code(output)
        global token_usage
        cur_prompt, cur_completion, cur_total = token_usage
        token_usage = (cur_prompt + usage[0], cur_completion + usage[1], cur_total + usage[2])
        logging.info("LLM输出：")
        logging.info(output)
        if (output == ""):
            logging.info("返回结果为空，下一轮优化")
            continue
        logging.info("替换并测试")
        function_replace(file_path, origin_code=origin_code, new_code=output)
        res, error_msg = cargo_test(test_target=test_target, multi_thread=multi_thread)
        if (res):
            cur_assert_count = output.count("assert!") + output.count("assert_eq!") + output.count("assert_ne!")
            logging.info(f"assert_count: {assert_count},   cur_assert_count:{cur_assert_count}")
            if cur_assert_count == assert_count:
                flag=True
                break
            else:
                logging.info(f"assert数量错误: 原{assert_count}, 现{cur_assert_count}")
                error_msg="The number of `assert` statements in the code you generated does not match that in the original code. Replace one `__assert_fail()` function with one macro, and do not change the quantity.If there are no `__assert_fail()` functions, do not add any assert macro."

        max_error_len = 1200
        logging.info(error_msg[:max_error_len])
        logging.info(f"cnt={cnt} 测试失败，还原重试")
        function_restore(file_path, origin_code=origin_code, new_code=output)
        cnt += 1
        # 补充错误信息
        if opt_round == 2 and "process didn't exit successfully" in error_msg:
            cur_user_ask = user_ask + fix_prompt_memory.format(last_code=output)
        else:
            cur_user_ask = user_ask + fix_prompt.format(last_code=output, error_msg=error_msg[:max_error_len])
            if "is defined multiple times" in error_msg:
                cur_user_ask += "Only output the Rust function, do not add any context, including definitions and use crates."
    if flag:
        return True
    else:
        return False

def incremental_optimize(tmp_file_path, user_ask, agent, test_target, assert_count, example_input=[], example_output=[], is_test=True):
    cnt = 0
    flag = False
    cur_user_ask = user_ask

    while cnt < MAX_ATTEMPS:
        logging.info("prompt输入：")
        logging.info(cur_user_ask)

        # 并行优化使用豆包
        output, usage = agent.generate_response(user_input=cur_user_ask, example_input=example_input, example_output=example_output)
        output = extract_rust_code(output)
        global token_usage
        cur_prompt, cur_completion, cur_total = token_usage
        token_usage = (cur_prompt + usage[0], cur_completion + usage[1], cur_total + usage[2])
        logging.info("LLM输出：")
        logging.info(output)
        if (output == ""):
            logging.info("返回结果为空，下一轮优化")
            continue

        logging.info("插入文件测试")
        target_content = output + "\n"
        if is_test:
            if "fn test_rb_tree_remove()" in output or "fn test_rb_tree_to_array()" in output:
                target_content = "#[test]\n#[ignore]\n" + output + "\n"
            else:
                target_content = "#[test]\n" + output + "\n"
        append_content(tmp_file_path, target_content)
        res, error_msg = cargo_test(test_target=test_target, multi_thread=True)
        if (res):
            cur_assert_count = output.count("assert!") + output.count("assert_eq!") + output.count("assert_ne!")
            logging.info(f"assert_count: {assert_count},   cur_assert_count:{cur_assert_count}")
            if cur_assert_count == assert_count:
                flag=True
                break
            else:
                logging.info(f"assert数量错误: 原{assert_count}, 现{cur_assert_count}")
                error_msg="The number of `assert` statements in the code you generated does not match that in the original code. Replace one `__assert_fail()` function with one macro, and do not change the quantity.If there are no `__assert_fail()` functions, do not add any assert macro."

            flag = True
            break
        max_error_len = 1200
        logging.info(error_msg[:max_error_len])
        logging.info(f"cnt={cnt} 测试失败，还原重试")
        remove_content(tmp_file_path, target_content)
        cnt += 1
        
        cur_user_ask = user_ask + incre_fix_prompt.format(error_msg=error_msg[:max_error_len])

            
    if flag:
        return True
    else:
        return False




# 遍历文件中的函数，进行assert宏优化
def assert_optimize(file_path, target=None):
    functions = split_rust_functions_in_file(file_path)
    with open(file_path,'r') as file:
        origin_content = file.read()
    for func in functions:
        if "s_test" in func['name']:
            continue
        origin_code = func['code']
        if origin_code not in origin_content:
            logging.info(f"下列原函数未能在文件{file_path}中查找到：\n{origin_code}")
            continue

        # example_in_1 = assert_optimize_prompt.format(rust_code = assert_example_code_1)
        # example_input = [
        #     example_in_1
        # ]
        # example_output = [
        #     assert_example_output_1
        # ]

        assert_count = origin_code.count("__assert_fail")

        user_ask = assert_optimize_prompt.format(rust_code=origin_code)
        
        # res = optimize(file_path, origin_code=origin_code, user_ask=user_ask, agent = doubao_agent, opt_round=0, assert_count=assert_count, test_target=target)
        # if not res:
        #     res_2 = optimize(file_path, origin_code=origin_code, user_ask=user_ask, agent = deepseek_agent, opt_round=0, assert_count=assert_count,  test_target=target)
        #     if not res_2:
        #         logging.info(f"函数{func['name']}assert优化失败")
        res = optimize(file_path, origin_code=origin_code, user_ask=user_ask, agent = doubao_agent, opt_round=0, assert_count=assert_count, test_target=target)
        if not res:
            if "test_tinyexpr" in file_path:
                agent = doubao_agent
            else:
                agent = deepseek_agent
            res_2 = optimize(file_path, origin_code=origin_code, user_ask=user_ask, agent = agent, opt_round=0, assert_count=assert_count,  test_target=target)
            if not res_2:
                logging.info(f"函数{func['name']}assert优化失败")
    



# 遍历文件中的函数，进行第一轮安全优化
# 弃用    
# def first_safe_optimize(file_path):
#     functions = split_rust_functions_in_file(file_path)
#     fail_list = []
#     with open(file_path,'r') as file:
#         origin_content = file.read()
#     example_input = [
#         first_optimize_prompt.format(rust_code=src_codein_1),
        
#     ]
#     example_output = [
#         src_codeout_1,
#     ]
#     for func in functions:
#         if "s_test" in func['name'] or "out_of_memory" in func['name']:
#             continue
#         origin_code = func['code']
#         if origin_code not in origin_content:
#             logging.info(f"下列原函数未能在文件{file_path}中查找到：\n{origin_code}")
#             continue
        
#         user_ask = first_optimize_prompt.format(rust_code=origin_code)
#         res = optimize(file_path, origin_code=origin_code, user_ask=user_ask, example_input=example_input, example_output=example_output, opt_round=1)
#         if not res:
#             logging.info(f"函数{func['name']}第一轮优化失败")
#             fail_list.append(func['name'])
#     return fail_list

def second_safe_optimize(file_path, target=None):
    functions = split_rust_functions_in_file(file_path)
    with open(file_path,'r') as file:
        origin_content = file.read()

    # 获取静态变量
    res = extract_static_declaration(path=file_path)
    static_code = '\n'.join(res) if res else ""

    # 获取extern C部分，利用extract_rust_code去注释
    res = extract_extern_declaration(path=file_path)
    extern_code = extract_rust_code(res if res else "")

    if "test_" in file_path:
        # logging.info("test file")
        example_in_1 = second_optimize_prompt.format(rust_code=test_second_opt_example_code_1, 
                                                static_variables=test_second_opt_example_static_1, 
                                                extern_C=test_second_opt_example_extern_C_1)
        example_in_2 = second_optimize_prompt.format(rust_code=test_second_opt_example_code_2, 
                                                static_variables=test_second_opt_example_static_2, 
                                                extern_C=test_second_opt_example_extern_C_2)
        example_input = [
            example_in_1, example_in_2
        ]
        example_output = [
            test_second_opt_example_output_1,test_second_opt_example_output_2
        ]
    else:
        example_in_1 = second_optimize_prompt.format(rust_code=src_second_opt_example_code_1, 
                                                static_variables=src_second_opt_example_static_1, 
                                                extern_C=src_second_opt_example_extern_C_1)
        example_in_2 = second_optimize_prompt.format(rust_code=src_second_opt_example_code_2, 
                                                static_variables=src_second_opt_example_static_2, 
                                                extern_C=src_second_opt_example_extern_C_2)
        example_in_3 = second_optimize_prompt.format(rust_code=src_second_opt_example_code_3, 
                                                static_variables=src_second_opt_example_static_3, 
                                                extern_C=src_second_opt_example_extern_C_3)
        example_input = [
            example_in_1, example_in_2, example_in_3
        ]
        example_output = [
            src_second_opt_example_output_1, src_second_opt_example_output_2, src_second_opt_example_output_3
        ]
    for func in functions:
        # 过滤
        if "test_functions" in func['name'] or "test_results" in func['name']:
            continue

        origin_code = func['code']
        if origin_code not in origin_content:
            logging.info(f"下列原函数未能在文件{file_path}中查找到：\n{origin_code}")
            continue
        user_ask = second_optimize_prompt.format(rust_code=origin_code, static_variables=static_code, extern_C=extern_code)
        assert_count = func['code'].count("assert!") + func['code'].count("assert_eq!") + func['code'].count("assert_ne!")
        # res = optimize(file_path, origin_code=origin_code, user_ask=user_ask, agent=deepseek_agent, example_input=example_input, example_output=example_output, opt_round=2, assert_count=assert_count, test_target=target)
        # if not res:
        #     logging.info(f"函数{func['name']}第二轮优化失败")
        res = optimize(file_path, origin_code=origin_code, user_ask=user_ask, agent=doubao_agent, example_input=example_input, example_output=example_output, opt_round=2, assert_count=assert_count, test_target=target)
        if not res:
            res_2 = optimize(file_path, origin_code=origin_code, user_ask=user_ask, agent=deepseek_agent, example_input=example_input, example_output=example_output, opt_round=2, assert_count=assert_count, test_target=target)
            if not res_2:
                logging.info(f"函数{func['name']}第二轮优化失败")
       

def global_variable_definition_optimize(file_path):
    logging.info("全局定义修改")
    res = extract_static_declaration(file_path)
    logging.info(res)
    if (not res) or (len(res) == 0):
        return ""
    # 过滤掉INIT_ARRAY和num_assert
    statics = [r for r in res if "INIT_ARRAY" not in r and "num_assert" not in r]
    static_code = '\n'.join(statics)
    logging.info(static_code)
    if static_code == "":
        return ""

    with open(file_path, 'r') as file:
        origin_content = file.read()
    new_content = origin_content
    insert_pos = new_content.find(statics[0])
    for s in statics:
        new_content = new_content.replace(f"{s}\n", "")
    
    example_in_1 = modify_variable_definition_prompt.format(static_variables=def_example_static_1)
    example_input = [ example_in_1 ]
    example_output = [ def_example_output_1 ]
    user_ask = modify_variable_definition_prompt.format(static_variables=static_code)

    logging.info("prompt输入：")
    logging.info(user_ask)
    if "utf8d" in static_code:
        output = tmp_definition
    else:
        output, usage = doubao_agent.generate_response(user_input=user_ask, example_input=example_input, example_output=example_output)
        output = extract_rust_code(output)
        global token_usage
        cur_prompt, cur_completion, cur_total = token_usage
        token_usage = (cur_prompt + usage[0], cur_completion + usage[1], cur_total + usage[2])
    logging.info("LLM输出：")
    logging.info(output)
    
    new_content = new_content[:insert_pos] + output + "\n" +  new_content[insert_pos:]
    with open(file_path, 'w') as file:
        file.write(new_content)
    return output

def test_parallel_optimize(file_path):
    file_name_without_ext = os.path.basename(file_path).split('.')[0]
    
    if "test_" in file_name_without_ext:
        test_module_name = file_name_without_ext
    # 此处未考虑 hash functions 和 compare functions
    else:
        test_module_name = f"test_{file_name_without_ext}"

    file_path_without_ext = os.path.splitext(file_path)[0]
    tmp_file_path = f"{file_path_without_ext}_tmp.rs"
    success_file_path = f"{file_path_without_ext}_success.rs"
    failed_file_path = f"{file_path_without_ext}_failed.rs"

    tmp_test_module_name = f"{test_module_name}_tmp"

    test_target = tmp_test_module_name if "test_" in file_name_without_ext else test_module_name
    print(f"test_target: {test_target}") 
    # 复制测试文件作为临时文件
    functions = split_rust_functions_in_file(file_path)
    shutil.copy(file_path, tmp_file_path)
    static_code = global_variable_definition_optimize(tmp_file_path)
    if static_code == "":
        os.remove(tmp_file_path)
        return

    with open(tmp_file_path, 'r') as file:
        origin_content = file.read()
    for func in functions:
        if "print_num_assert" in func['name']:
            continue

        if "test_" in func['name'] and "alloc_test" not in func['name']:
            if "test_rb_tree_to_array" in func['name'] or "test_rb_tree_remove" in func['name']:
                replace_target = "#[test]\n#[ignore]\n"+func['code']
            else:
                replace_target = "#[test]\n"+func['code']
        else:
            replace_target = func['code']
        origin_content = origin_content.replace(replace_target, "")
        origin_content = origin_content.replace("#[inline]", "")

    with open(tmp_file_path, 'w') as file:
        file.write(origin_content)

    # 构造prompt
    example_in_1 = parallel_opt_prompt.format(rust_code=pal_opt_example_code_1, 
                                        static_variables=pal_opt_example_def_1)
    example_in_2 = parallel_opt_prompt.format(rust_code=pal_opt_example_code_2, 
                                            static_variables=pal_opt_example_def_2)
    example_in_3 = parallel_opt_prompt.format(rust_code=pal_opt_example_code_3, 
                                            static_variables=pal_opt_example_def_3)
    example_input = [
        example_in_1, example_in_2,  example_in_3
    ]
    example_output = [
        pal_opt_example_output_1, pal_opt_example_output_2, pal_opt_example_output_3, 
    ]

    for func in functions:
        if "test_" in func['name'] and "alloc_test" not in func['name']:
            is_test = True
        else:
            is_test = False
        user_ask = parallel_opt_prompt.format(rust_code=func['code'], static_variables=static_code)
        assert_count = func['code'].count("assert!") + func['code'].count("assert_eq!") + func['code'].count("assert_ne!")
        
        res = incremental_optimize(tmp_file_path, user_ask, agent = doubao_agent, test_target=test_target, assert_count=assert_count, example_input=example_input, example_output=example_output, is_test=is_test)

        if not res:
            res_2 = incremental_optimize(tmp_file_path, user_ask, agent = deepseek_agent, test_target=test_target, assert_count=assert_count, example_input=example_input, example_output=example_output, is_test=is_test)
            if not res_2:
                logging.info(f"文件{os.path.basename(file_path)}并行性优化失败")
                os.remove(tmp_file_path)
                # os.rename(file_path, failed_file_path)
                # os.rename(tmp_file_path, file_path)
                return False
    logging.info(f"文件{os.path.basename(file_path)}并行性优化成功")
    os.replace(tmp_file_path, file_path)
    return True      
    # os.rename(file_path, success_file_path)
    # os.rename(tmp_file_path, file_path)



def src_parallel_optimize(file_path):
    file_path_without_ext = os.path.splitext(file_path)[0]
    tmp_file_path = f"{file_path_without_ext}_tmp.rs"
    success_file_path = f"{file_path_without_ext}_success.rs"
    failed_file_path = f"{file_path_without_ext}_failed.rs"

    file_name_without_ext = os.path.basename(file_path).split('.')[0]
    test_target = f"test_{file_name_without_ext}"
    print(f"test_target: {test_target}")

    # 备份
    shutil.copy(file_path, tmp_file_path)
    functions = split_rust_functions_in_file(file_path)
    static_code = global_variable_definition_optimize(file_path)
    if static_code == "":
        return
    
    with open(file_path, 'r') as file:
        origin_content = file.read()

    # 构造prompt
    example_in_1 = parallel_opt_prompt.format(rust_code=src_pal_opt_example_code_1, 
                                        static_variables=src_pal_opt_example_def_1)
    example_in_2 = parallel_opt_prompt.format(rust_code=src_pal_opt_example_code_2, 
                                            static_variables=src_pal_opt_example_def_2)

    example_input = [
        example_in_1, example_in_2
    ]
    example_output = [
        src_pal_opt_example_output_1, src_pal_opt_example_output_2
    ]


    cnt = 0
    flag = False
    while cnt < MAX_ATTEMPS:
        new_content = origin_content
        for func in functions:
            user_ask = parallel_opt_prompt.format(rust_code=func['code'], static_variables=static_code)
            output, usage = doubao_agent.generate_response(user_input=user_ask, example_input=example_input, example_output=example_output)
            output = extract_rust_code(output)
            global token_usage
            cur_prompt, cur_completion, cur_total = token_usage
            token_usage = (cur_prompt + usage[0], cur_completion + usage[1], cur_total + usage[2])
            new_content = new_content.replace(func['code'], output)
        with open(file_path, 'w') as file:
            file.write(new_content)
        res, error_msg = cargo_test(test_target=test_target, multi_thread=True)
        if (res):
            flag = True
            break
        logging.info(error_msg)
        logging.info(f"cnt={cnt} 测试失败，重试")
        cnt += 1
        shutil.copy(tmp_file_path, file_path)
        static_code = global_variable_definition_optimize(file_path)

    if flag:
        logging.info(f"文件{os.path.basename(file_path)}并行性优化成功！")
        # os.rename(file_path, success_file_path)
        # os.rename(tmp_file_path, file_path)
        os.remove(tmp_file_path)
    else:
        cnt = 0
        flag = False
        while cnt < MAX_ATTEMPS:
            new_content = origin_content
            for func in functions:
                user_ask = parallel_opt_prompt.format(rust_code=func['code'], static_variables=static_code)
                output, usage = deepseek_agent.generate_response(user_input=user_ask, example_input=example_input, example_output=example_output)
                output = extract_rust_code(output)
                # global token_usage
                # cur_prompt, cur_completion, cur_total = token_usage
                # token_usage = (cur_prompt + usage[0], cur_completion + usage[1], cur_total + usage[2])
                new_content = new_content.replace(func['code'], output)
            with open(file_path, 'w') as file:
                file.write(new_content)
            res, error_msg = cargo_test(test_target=test_target, multi_thread=True)
            if (res):
                flag = True
                break
            logging.info(error_msg)
            logging.info(f"cnt={cnt} 测试失败，重试")
            cnt += 1
            shutil.copy(tmp_file_path, file_path)
            static_code = global_variable_definition_optimize(file_path)

        if flag:
            logging.info(f"文件{os.path.basename(file_path)}并行性优化成功！")
            os.remove(tmp_file_path)
        else:
            logging.info(f"文件{os.path.basename(file_path)}并行性优化失败")
            # os.rename(file_path, failed_file_path)
            os.replace(tmp_file_path, file_path)
    

def rewrite_assert_macro(file_path):
    with open(file_path , 'r') as file:
        lines = file.readlines()
    new_lines = [l for l in lines if "num_assert = num_assert.wrapping_add" not in l]
    new_lines = [l for l in new_lines if "pub static mut num_assert" not in l]
    with open(file_path, 'w') as file:
        file.writelines(new_lines)

    with open(file_path , 'r') as file:
        origin_content = file.read()
    if new_macro not in origin_content:
        origin_content += new_macro
    with open(file_path, 'w') as file:
        file.write(origin_content)

print_function="""
#[dtor]
fn print_num_assert() {
    println!("num_assert: {}", num_assert.load(Ordering::Relaxed))
}
"""
def add_print_function(file_path):
    with open(file_path, 'a') as file:
        file.write(print_function)

# # 移除（注释）test 文件中的extern C函数声明
# # fail_list中为未能转为safe的函数名，保留其extern C声明
# def remove_extern_declaration_for_testfile(file_path):
#     # 文件名 test_avl_tree.rs
#     file_name = os.path.basename(file_path)
#     # 提取模块名 avl_tree
#     start_index = len("test_")
#     end_index = file_name.find(".rs")
#     module_name = file_name[start_index:end_index]

#     remove_list = set()
#     for file in COMPARE_HASH_SRC_LIST:
#         res = extract_function_names(file)
#         print(res)
#         if res:
#             remove_list = remove_list | set(res)
#     if "hash_functions" not in module_name and "compare_functions" not in module_name:
#         res = extract_function_names(f"{SRC_PATH}/{module_name}.rs")
#         if res:
#             remove_list = remove_list | set(res)
#     # remove_list = remove_list - set(fail_list)

#     res = extract_extern_declaration(file_path)
#     extern_decl = res if res else ""
#     func_name_to_sig = extract_func_signature_from_extern_declaration(extern_decl)
#     if func_name_to_sig:
#         for func_name in func_name_to_sig.keys():
#             if func_name in remove_list or "run_tests" in func_name or "alloc_test" in func_name:
#                 comment_code(file_path=file_path, target_code=func_name_to_sig[func_name])

# def remove_extern_declaration_for_srcfile(file_path):
#     res = extract_extern_declaration(file_path)
#     extern_decl = res if res else ""

#     func_name_to_sig = extract_func_signature_from_extern_declaration(extern_decl)
#     if func_name_to_sig:
#         for func_name in func_name_to_sig.keys():
#             if "alloc_test" in func_name:
#                 comment_code(file_path=file_path, target_code=func_name_to_sig[func_name])




# # 移除（注释）test 文件中的结构体声明
# def remove_struct_and_type_declaration(file_path):
#     res = extract_struct_declaration(file_path)
#     struct_codes = res if res else []

#     logging.info(f"查找到{len(struct_codes)}个结构体定义")
#     for struct in struct_codes:
#         comment_code(file_path, struct)
#     attribute="#[derive(Copy, Clone)]\n#[repr(C)]"
#     comment_code(file_path, attribute)

#     res = extract_type_declaration(file_path)
#     type_codes = res if res else []
#     logging.info(f"查找到{len(type_codes)}个type定义")
#     for type in type_codes:
#         comment_code(file_path, type)


# 从c2rust的翻译结果构造rust项目
def construct_project():
    if os.path.exists(f"{C2RUST_OUTPUT_PATH}/src/src") and os.path.exists(f"{C2RUST_OUTPUT_PATH}/src/test"):
        if not os.path.exists(f"{PROJ_PATH}"):
            logging.info("项目文件夹不存在，新建cargo项目")
            result = subprocess.run(
                ["cargo", "new", f"{PROJ_PATH}"],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
            )
            with open(f"{PROJ_PATH}/Cargo.toml", 'a') as file:
                file.write('libc = "0.2.167"\n')
                file.write('lazy_static = "1.4.0"\n')
                file.write('\n[dev-dependencies]\n')
                file.write('ctor = "0.1"\n')
                
            shutil.rmtree(f"{PROJ_PATH}/src")
        else:
            logging.info("项目文件夹存在，移除src、tests目录")
            if os.path.exists(f"{PROJ_PATH}/src"):
                shutil.rmtree(f"{PROJ_PATH}/src")
            if os.path.exists(f"{PROJ_PATH}/tests"):
                shutil.rmtree(f"{PROJ_PATH}/tests")

        shutil.copytree(f"{C2RUST_OUTPUT_PATH}/src/src", f"{PROJ_PATH}/src")
        shutil.copytree(f"{C2RUST_OUTPUT_PATH}/src/test", f"{PROJ_PATH}/tests")
    os.remove(f"{PROJ_PATH}/tests/framework.rs")
    shutil.move(f"{PROJ_PATH}/tests/alloc_testing.rs", f"{PROJ_PATH}/src/alloc_testing.rs")
    if not os.path.exists(f"{PROJ_PATH}/count_safe_ratio.sh"):
        shutil.copy(COUNT_SAFE_RATIO_PATH, f"{PROJ_PATH}")
        subprocess.run(['chmod', '+x', f"{PROJ_PATH}/count_safe_ratio.sh"])
    add_crate_and_lib(PROJ_PATH)
    # remove_out_of_memory_test(PROJ_PATH)
    add_alloc_set_limit(PROJ_PATH)
    remove_main_and_tests(PROJ_PATH)
    remove_c_label_blocks(PROJ_PATH)
    add_stest_functions(PROJ_PATH)
    
def make_c_project():
    test_cpp_path = f"{C_PATH}/test/test-cpp.cpp"
    test_arraylist_path = f"{C_PATH}/test/test-arraylist.c"
    cpp_assert_macro = extract_assert_macro(test_cpp_path)
    c_assert_macro = extract_assert_macro(test_arraylist_path)
    test_cpp_path_c = f"{C_PATH}/test/test-cpp.c"
    shutil.copy(test_cpp_path, test_cpp_path_c)
    function_replace(test_cpp_path_c, cpp_assert_macro, c_assert_macro)

    current_directory = os.getcwd()
    os.chdir(f"{C_PATH}")
    result = subprocess.run(
        ["make", "clean"],
        stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        text=True,
    )
    result = subprocess.run(
        ["bear", "--", "make", "all"],
        stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        text=True,
    )
    os.chdir(current_directory)
    if result.returncode == 0:
        return True
    else:
        return False

def c2rust_convert():
    if os.path.exists(C2RUST_OUTPUT_PATH):
        logging.info("c2rust输出文件夹存在，清空")
        shutil.rmtree(C2RUST_OUTPUT_PATH)
    os.makedirs(C2RUST_OUTPUT_PATH)
    result = subprocess.run(
        [C2RUST_PATH, "transpile", "--output-dir", f"{C2RUST_OUTPUT_PATH}", f"{C_PATH}/compile_commands.json"],
        stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        text=True,
    )
    if result.returncode == 0:
        return True
    else:
        return False

# rustfmt格式化
def rustfmt_format():
    current_directory = os.getcwd()
    os.chdir(f"{PROJ_PATH}")
    # 第一次构建，可能有crates更新   
    result = subprocess.run(
            ["cargo", "fmt"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
    os.chdir(current_directory)

    if result.returncode == 0:
        return True
    else:
        return False
        

if __name__ == "__main__":
    # 日志
    setup_logging()    
    logging.info(f"项目{CRATE_NAME}")
    logging.info(f"路径{PROJ_PATH}")

    # 编译命令文件
    logging.info("编译C代码")
    if not make_c_project():
        logging.info("C代码编译失败")
        exit(1)
    logging.info("C代码编译成功")

    # C代码预处理
    logging.info("预处理C代码")
    # remove_alloc_test_blocks_in_dir(f"{C_PATH}/test")
    remove_atomic_type(f"{C_PATH}/test")


    # C2Rust
    logging.info("c2rust转换")
    c2rust_start_time = time.time()
    
    res = c2rust_convert()
    if res == True:
        logging.info("c2rust 转换成功")
    else:
        logging.info("c2rust 转换失败")
        exit(2)
        
    # 构建项目并测试
    logging.info("项目构建")
    construct_project()

    logging.info("c2rust 测试")
    flag, err = cargo_test(FIRST_TEST_TIMEOUT)
    c2rust_end_time = time.time()
    logging.info(f"c2rust转换耗时{str(datetime.timedelta(seconds=(c2rust_end_time - c2rust_start_time)))}")
    if not flag:
        exit(1)


    # rustfmt格式化，第一轮优化，再次格式化
    # 移除（注释）文件中的extern C函数声明
    remove_extern_declaration_for_srcfile(PROJ_PATH)
    remove_extern_declaration_for_testfile(PROJ_PATH)
    # 移除（注释）test 文件中与src文件重复的结构体声明
    remove_struct_and_type_declaration(PROJ_PATH)
    # 移除stest函数
    remove_stest_functions(PROJ_PATH)
    # 明确指明size_t来源
    add_use_crate(PROJ_PATH, f"use {CRATE_NAME}::alloc_testing::size_t;\n")
    # rustfmt格式化
    rustfmt_format()

    logging.info("第一轮优化")
    first_start_time = time.time()

    first_round_optimize(PROJ_PATH)
    rustfmt_format()
    # 为原来的测试函数添加测试注解
    add_test_annotation(PROJ_PATH)
    # 移除no_mangle注解
    remove_no_mangle_annotation(PROJ_PATH)
    # 运行测试
    cargo_test(TEST_TIMEOUT)

    first_end_time = time.time()
    logging.info(f"第一轮优化完成，耗时{str(datetime.timedelta(seconds=(first_end_time - first_start_time)))}")
    
    # 调试用
    # DEBUG_FLAG=True
    DEBUG_FLAG = False
    DEBUG_TARGET=[
                # "test_cpp.rs",

                "alloc_testing.rs", 
                # "test_alloc_testing.rs",
                
                # "arraylist.rs", 
                # "test_arraylist.rs",

                # "avl_tree.rs", 
                # "test_avl_tree.rs",

                # "binary_heap.rs",
                # "test_binary_heap.rs",

                # "binomial_heap.rs",
                # "test_binomial_heap.rs",

                # "bloom_filter.rs",
                # "test_bloom_filter.rs"


                # "compare_int.rs",
                # "compare_pointer.rs",
                # "compare_string.rs",
                # "test_compare_functions.rs",

                # "hash_int.rs",
                # "hash_pointer.rs",
                # "hash_string.rs",
                # "test_hash_functions.rs",

                # "hash_table.rs",
                # "test_hash_table.rs",

                # "list.rs",
                # "test_list.rs",

                # "queue.rs",
                # "test_queue.rs",

                # "rb_tree.rs", 
                # "test_rb_tree.rs",

                
                # "set.rs", 
                # "test_set.rs",

                # "slist.rs", 
                "test_slist.rs",

                # "sortedarray.rs",
                # "test_sortedarray.rs",

                # "tinyexpr.rs", 
                "test_tinyexpr.rs",
                
                # "trie.rs",
                # "test_trie.rs",
                
                # "utf8_decoder.rs", 
                "test_utf8_decoder.rs",
                ]

    # assert宏替换
    assert_start_time = time.time()
    for root, dirs, files in os.walk(TEST_PATH):
        files.sort()
        for file in files:
            if DEBUG_FLAG and file not in DEBUG_TARGET:
                continue
            logging.info(f"注释{file}中的main函数")
            logging.info("assert宏替换")

            target = file.split('.')[0]
            assert_optimize(os.path.join(root, file), target = target)


    for root, dirs, files in os.walk(SRC_PATH):
        files.sort()
        for file in files:
            if  DEBUG_FLAG and file not in DEBUG_TARGET:
                continue
            if "alloc_testing" in file:
                logging.info("当前文件为alloc_testing，先进行assert替换")

                assert_optimize(file_path = os.path.join(root, file))

    assert_end_time = time.time()
    logging.info(f"assert宏替换完成，耗时{str(datetime.timedelta(seconds=(assert_end_time - assert_start_time)))}")
    logging.info(f"token 消耗统计：")
    print(f"Prompt Tokens: {token_usage[0]}")
    print(f"Completion Tokens: {token_usage[1]}")
    print(f"Total Tokens: {token_usage[2]}")



    # 宏重写
    for root, dirs, files in os.walk(TEST_PATH):
        files.sort()
        for file in files:
            if DEBUG_FLAG and file not in DEBUG_TARGET:
                continue
            file_path = os.path.join(root, file)

            logging.info(f"{file} 宏重写")
            rewrite_assert_macro(file_path)
    cargo_test(TEST_TIMEOUT)


    # 第二轮优化
    second_start_time = time.time()
    for root, dirs, files in os.walk(SRC_PATH):
        files.sort()
        for file in files:
            if  DEBUG_FLAG and file not in DEBUG_TARGET:
                continue

            module_name = file.split('.')[0]
            file_path = os.path.join(root, file)

            remove_extern_declaration_for_srcfile(file_path)
            logging.info(f"{file}第二轮优化")

            if "compare_int" in file or "compare_string" in file or "compare_pointer" in file:
                target = "test_compare_functions"
            elif "hash_int" in file or "hash_string" in file or "hash_pointer" in file:
                target = "test_hash_functions"
            else:
                target = f"test_{module_name}"
            second_safe_optimize(file_path, target)


    for root, dirs, files in os.walk(TEST_PATH):
        files.sort()
        for file in files:
            if  DEBUG_FLAG and file not in DEBUG_TARGET:
                continue
            
            start_index = len("test_")
            end_index = file.find(".rs")
            module_name = file[start_index:end_index]

            file_path = os.path.join(root, file)
           
            logging.info(f"{file}第二轮优化")
            target = file.split('.')[0]
            second_safe_optimize(file_path, target=target)
            
    second_end_time = time.time()
    logging.info(f"第二轮优化完成，耗时{str(datetime.timedelta(seconds=(second_end_time - second_start_time)))}")
    logging.info(f"token 消耗统计：")
    print(f"Prompt Tokens: {token_usage[0]}")
    print(f"Completion Tokens: {token_usage[1]}")
    print(f"Total Tokens: {token_usage[2]}")




    # 并行性优化
    parallel_start_time = time.time()
    for root, dirs, files in os.walk(SRC_PATH):
        files.sort
        for file in files:
            if  DEBUG_FLAG and file not in DEBUG_TARGET:
                continue
            logging.info(f"{file}并行性优化")
            if "alloc_testing" in file:
                src_parallel_optimize(os.path.join(root, file))

    for root, dirs, files in os.walk(TEST_PATH):
        files.sort()
        for file in files:
            if  DEBUG_FLAG and file not in DEBUG_TARGET:
                continue
            file_path = os.path.join(root, file)
            logging.info(f"{file}并行性优化")
            res = test_parallel_optimize(file_path)
    parallel_end_time = time.time()
    logging.info(f"并行性优化完成，耗时{str(datetime.timedelta(seconds=(parallel_end_time - parallel_start_time)))}")

    # 添加打印函数
    logging.info(f"添加打印函数")
    for root, dirs, files in os.walk(TEST_PATH):
        files.sort()
        for file in files:
            if  DEBUG_FLAG and file not in DEBUG_TARGET:
                continue
            file_path = os.path.join(root, file)
            add_print_function(file_path)


    logging.info(f"token 消耗统计：")
    print(f"Prompt Tokens: {token_usage[0]}")
    print(f"Completion Tokens: {token_usage[1]}")
    print(f"Total Tokens: {token_usage[2]}")
    

    

    # # safe 统计
    # logging.info("统计safe指标")
    # current_directory = os.getcwd()
    # os.chdir(f"{PROJ_PATH}")
    # process = subprocess.Popen("./count_safe_ratio.sh", stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    # stdout, _ = process.communicate()
    # logging.info(f"{stdout}")
    # os.chdir(current_directory)

    
    
